/*
 * Copyright 2012 Alain Pannetrat (L1L1@gmx.com)
 * 
 * This file is part of cookie-miner.
 * 
 * cookie-miner is free software: you can redistribute it and/or modify 
 * it under the terms of the GNU General Public License as published by 
 * the Free Software Foundation, either version 3 of the License, or 
 * (at your option) any later version.
 * 
 * cookie-miner is distributed in the hope that it will be useful, but 
 * WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU 
 * General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License 
 * along with cookie-miner. If not, see http://www.gnu.org/licenses/.
 *
 * @version: $Id$
 *
 */
#include "http_proxy.h"
#include "bufio.h"
#include "bufio_ssl.h"
#include "http_packet.h"
#include "log_facility.h"
#include "ssl_tools.h"
#include "ssl_tools.h"
#include "tcp.h"
#include <pthread.h>
#include <string>
#include <sys/stat.h>
#include <vector>
#include "url.h"

class HttpProxyInstance : private Uncopyable {
	public:

		static HttpProxyInstance *create(const HttpProxy* proxy, int socket);
		bool run();
		~HttpProxyInstance();

		enum proxy_state {
		        PROXY_INIT,
			PROXY_REQUEST,
			PROXY_RESPONSE,
			PROXY_TUNNEL,
			PROXY_SSL_INIT,
			PROXY_SSL_REQUEST,
			PROXY_SSL_RESPONSE,
			PROXY_RESPONSE_GATEWAY_TIMEOUT,
			PROXY_RESPONSE_NOT_FOUND,
			PROXY_CONNECTION_CLOSE,
			PROXY_RESPONSE_INTERNAL,
			PROXY_FILTER_BLOCK,
			PROXY_ERROR
		};

	private:
		bool ssl_init();
		
		HttpProxyInstance(const HttpProxy* proxy, int socket, proxy_state state = PROXY_INIT);

		const HttpProxy* _proxy;
		BufferInOutSocket *_client;
		BufferInOutSocket *_server;
		HttpPacket *_previous_packet;
		HttpPacket *_current_packet;
		bool _save_packet;
		unsigned _content_length_limit;

		bool has_prev_packet() { return _previous_packet!=NULL; }
		bool has_packet() { return _current_packet!=NULL; }
		HttpPacket& packet() { return *_current_packet; }
		HttpPacket& prev_packet() { return *_previous_packet; }
		HttpPacket& new_packet();

                proxy_state process_request(void);
		proxy_state process_response_gateway_timeout(void);
		proxy_state process_response_not_found(void);
		proxy_state process_response(void);
		proxy_state process_tunnel_setup(void);
		proxy_state process_ssl_request(void);
		proxy_state process_ssl_response(void);
		proxy_state process_response_internal(void);

		EVP_PKEY* local_client_key;
		EVP_PKEY* local_server_key;
		EVP_PKEY* local_ca_key;
		X509*	  local_ca_cert;
};

/********************************************************************
 *
 *	HttpProxy
 *
 */

typedef struct {	
	int 		thread_socket;
	HttpProxy* 	thread_proxy;
} thread_context_t;

void *proxy_instance_thread(void *ctx_ptr)
{
	thread_context_t *ctx = (thread_context_t *)ctx_ptr;
	HttpProxyInstance *proxy;

	logger.new_thread();
	logger.message(logger.LOG_DEBUG,"Creating proxy thread with socket %i.",ctx->thread_socket);

	proxy = HttpProxyInstance::create(ctx->thread_proxy, ctx->thread_socket);
	proxy->run();
	delete proxy;

	tcp_close(ctx->thread_socket);
	logger.message(logger.LOG_DEBUG,"Exiting proxy thread.");
	delete ctx;
	pthread_exit(NULL);
	return NULL;
}

HttpProxy::HttpProxy(Json::Value* config) 
{
	if (config==NULL) 
		_config = Json::create_object(); 
	else 
		_config = config; 
} 

bool HttpProxy::run()
{
	char ip[20];
	unsigned short port;
	int sock;
	thread_context_t *ctx;
	pthread_t thread_ref;

	if (_config->get("proxy")->get("port")->is_defined())
		port = _config->get("proxy")->get("port")->as_int();
	else
		port = 8080;

	if ((sock = tcp_listen(port))<0)
		return false;

	for (;;)
	{
		ctx = new thread_context_t;
		ctx->thread_proxy  			= this;
		ctx->thread_socket 			= tcp_accept(sock,ip);

		pthread_create(&thread_ref,NULL,&proxy_instance_thread,ctx);
		pthread_detach(thread_ref);
	}
}

bool HttpProxy::filter_packet(HttpPacket& packet) const
{
	HttpFilterList::const_iterator it;

	for (it=_filters.begin(); it!=_filters.end(); it++)
	{
		if (!(*it)->apply(packet)) return false;
	}
	return true;
}

HttpProxy::~HttpProxy()
{
	HttpFilterList::const_iterator it;

	for (it=_filters.begin(); it!=_filters.end(); it++)
	{
		delete (*it);
	}	
}

/********************************************************************
 /*
 *	HttpProxyInstance
 *
 */

HttpProxyInstance::HttpProxyInstance(const HttpProxy* proxy, int socket, proxy_state state) :
		        _proxy(proxy),	
			_client(BufferInOutSocket::create(socket)),
			_server(NULL), 
			_previous_packet(NULL),
       			_current_packet(NULL) 
{
	if (proxy->config()->get("proxy")->get("save_packets")->is_defined())
		_save_packet = proxy->config()->get("proxy")->get("save_packets")->as_bool();
	else
		_save_packet = true;

	if (proxy->config()->get("proxy")->get("content_length_limit")->is_defined())
		_content_length_limit = (unsigned)proxy->config()->get("proxy")->get("content_length_limit")->as_int();
	else
		_content_length_limit = (unsigned)(-1);
}

/*private*/ HttpPacket& HttpProxyInstance::new_packet() 
{ 
	if (_previous_packet) 
		delete _previous_packet; 

	_previous_packet = _current_packet;

	_current_packet = HttpPacket::create();

	_current_packet->content_length_limit(_content_length_limit);

	if (_previous_packet) 
		_current_packet->parent_id(_previous_packet->id());

	return *_current_packet; 
}

bool HttpProxyInstance::ssl_init()
{
    struct stat st;

//    if (local_ca_cert!=NULL)
//    	return true;

    if (stat("KEYS",&st)!=0)
    {
	logger.message(logger.LOG_WARNING,"Could not find ./KEYS subdirectory for SSL in current working directory.");
	return false;
    }

    if ((local_client_key = openssl_load_private_key_from_file("KEYS/client_key_pair.key","secret"))==NULL)
    	return false;
    if ((local_server_key = openssl_load_private_key_from_file("KEYS/server_key_pair.key","secret"))==NULL)
	return false;
    if ((local_ca_key = openssl_load_private_key_from_file("KEYS/ca_key_pair.key","secret"))==NULL)
    	return false;
    if ((local_ca_cert = openssl_load_cert_from_file("KEYS/ca_cert.pem"))==NULL)
	return false;
    return true;
}

HttpProxyInstance *HttpProxyInstance::create(const HttpProxy *proxy, int socket)
{
    return new HttpProxyInstance(proxy, socket);
}

HttpProxyInstance::~HttpProxyInstance()
{
    if (_client) delete _client; // close is called elsewhere in the thread
    
    if (_server) 
    {
	    tcp_close(_server->get_fd());
	    delete _server;
    }
    if (_previous_packet) delete _previous_packet;
}

// 
// PROXY_TUNNEL
//
HttpProxyInstance::proxy_state HttpProxyInstance::process_tunnel_setup()
{
	const std::string &host = packet().headline().part(1);
	X509*     local_server_cert;
	X509*     distant_server_cert;
	std::string description;
	X509*	  cert_chain[2];
	std::string line;

	new_packet();

	if (!ssl_init())
	{
		logger.message(logger.LOG_ERROR,"Failed to initilaize SSL");
		return PROXY_ERROR;
	}

	int sock = tcp_connect(host.c_str());

	if (sock<0)
	{	
		logger.message(logger.LOG_WARNING,"failed to establish connection to %s, sending error 504.",host.c_str());	
		return process_response_gateway_timeout();	
	}

	logger.message(logger.LOG_DEBUG,"Connected socket to %s",host.c_str());
	_server = BufferInOutSSL::create(sock,false,local_client_key);
	if (!_server)
	{
		return PROXY_ERROR; /* TODO: release SSL alloc objects */
	}

	logger.message(logger.LOG_DEBUG,"Established SSL client connection proxy<-->%s",host.c_str());

	packet().headline().set("HTTP/1.1 200 Connection established\r\n");
	packet().headers().set("Proxy-agent","Proxy Certs");

	packet().source(socket_get_local_address(_client->get_fd()));
	packet().destination(socket_get_peer_address(_client->get_fd()));

	if (!packet().headline().write_out(_client) ||
	    !packet().headers().write_out(_client) || 
	    !_client->write_flush())
	{
		return PROXY_ERROR; /* TODO: release SSL alloc objects */
	}

	packet().log();

	distant_server_cert = SSL_get_peer_certificate(((BufferInOutSSL*)_server)->get_ssl());
	
	if (distant_server_cert==NULL)
	{
		logger.message(logger.LOG_ERROR,"Failed to get peer certificate");
		return PROXY_ERROR;
	}
	local_server_cert = openssl_transform_certificate(distant_server_cert,local_ca_key,local_server_key);

	sock = _client->get_fd();
	delete _client;

	cert_chain[0]=local_server_cert;
	cert_chain[1]=local_ca_cert;

	logger.message(logger.LOG_DEBUG,"Upgrading browser<-->proxy connection to SSL");

	_client = BufferInOutSSL::create(sock, true, local_server_key, 2, cert_chain);
	if (!_client)
	{
		return PROXY_ERROR;
		/* TODO: cleanup */
	}

	logger.message(logger.LOG_DEBUG,"established SSL server connection from %s:%u to %s:%u",
			packet().source().ip(),
			packet().source().port(),
			packet().destination().ip(),
			packet().destination().port()
		      );

	return PROXY_SSL_REQUEST;
}

//
// PROXY_SSL_REQUEST
//
HttpProxyInstance::proxy_state HttpProxyInstance::process_ssl_request()
{
    new_packet();

    packet().source(socket_get_peer_address(_client->get_fd()));
    packet().destination(socket_get_peer_address(_server->get_fd()));

    if (!packet().headline().read_in(_client) || 
	!packet().headers().read_in(_client))
    {
	return PROXY_ERROR;
    }

    logger.message(logger.LOG_DEBUG,"Processing ssl request %s",packet().headline().full().c_str());

    packet().log();	

    if (!_proxy->filter_packet(packet())) return PROXY_FILTER_BLOCK; 

    if (!packet().headline().write_out(_server) ||
	!packet().headers().write_out(_server) ||
    	!packet().content_read_in_write_out(_client,_server) || 
	!_server->write_flush())
    {
	return PROXY_ERROR;
    }    
    
    return PROXY_SSL_RESPONSE;
}


//
// PROXY_REQUEST
//
HttpProxyInstance::proxy_state HttpProxyInstance::process_request()
{
    std::string host;

    new_packet(); 

    if (!packet().headline().read_in(_client) || 
	!packet().headers().read_in(_client))
    {
	return PROXY_ERROR;
    }

    logger.message(logger.LOG_DEBUG,"Processing request %s",packet().headline().full().c_str());

    const std::string& http_method 	= packet().headline().part(0);
    const std::string& http_ressource 	= packet().headline().part(1);
    const std::string& http_proto	= packet().headline().part(2);

    packet().log();	

    packet().source(socket_get_peer_address(_client->get_fd()));

    if (!_proxy->filter_packet(packet())) return PROXY_FILTER_BLOCK; 

    if (http_method=="CONNECT") /* assume SSL */
    {
	    packet().destination(socket_get_local_address(_client->get_fd()));

	    return PROXY_TUNNEL;
    }

    if (!packet().headers().exists("Host")) 
    {
	    logger.message(logger.LOG_ERROR,"Missing 'Host:' header in message");
	    return PROXY_ERROR;
    }
    
    URL url(http_ressource);

    if (url.domain()=="")
    {
	    packet().source(socket_get_peer_address(_client->get_fd()));
	    packet().destination(socket_get_local_address(_client->get_fd()));

	    return PROXY_RESPONSE_INTERNAL;
    }

    if (url.has_port())
    	host = url.domain() + ":" + url.port();
    else
	host = url.domain();

    logger.message(logger.LOG_DEBUG,"Realtive URL is %s",url.encode_relative().c_str());
    packet().headline().set(http_method,url.encode_relative(),http_proto);
    logger.message(logger.LOG_DEBUG,"Translated request to %s",packet().headline().full().c_str());

    if (_server==NULL) 
    {
	    int sock;

	    sock = tcp_connect(host.c_str());

	    if (sock<0) 
	    {
		    logger.message(logger.LOG_ERROR,"Failed to establish connection to %s",host.c_str());
		    packet().source(socket_get_peer_address(_client->get_fd()));
		    packet().destination(socket_get_local_address(_client->get_fd()));

		    return PROXY_RESPONSE_NOT_FOUND;
	    }

	    logger.message(logger.LOG_DEBUG,"established connection to %s",host.c_str());
	    _server = BufferInOutSocket::create(sock);
    }

    packet().destination(socket_get_peer_address(_server->get_fd()));

    if (!packet().headline().write_out(_server) ||
    	!packet().headers().write_out(_server) ||
    	!packet().content_read_in_write_out(_client,_server) || 
    	!_server->write_flush())
    {
	    logger.message(logger.LOG_ERROR,"Could not write packet to destination");
	    return PROXY_ERROR;
    }

    return PROXY_RESPONSE; 
}

//
// PROXY_SSL_RESPONSE
//
HttpProxyInstance::proxy_state HttpProxyInstance::process_ssl_response()
{
    new_packet(); 

    packet().source(socket_get_peer_address(_server->get_fd()));
    packet().destination(socket_get_peer_address(_client->get_fd()));

    if (!packet().headline().read_in(_server) || 
	!packet().headers().read_in(_server))
    {
	return PROXY_ERROR;
    }

    if (packet().headline().full().find("HTTP")!=0)
    {	
	logger.message(logger.LOG_ERROR,"Malformed SSL response header: %s",packet().headline().full().c_str());
	return PROXY_ERROR;
    }
    
    logger.message(logger.LOG_DEBUG,"Processing SSL response %s",packet().headline().full().c_str());
 
    packet().log();

    if (!_proxy->filter_packet(packet())) return PROXY_FILTER_BLOCK; 

    if (!packet().headline().write_out(_client) || 
        !packet().headers().write_out(_client) ||
    	!packet().content_read_in_write_out(_server,_client) ||
	!_client->write_flush())
    {
	logger.message(logger.LOG_ERROR,"Failed to write SSL response packet to client");
	return PROXY_ERROR; 
    }

    return PROXY_SSL_REQUEST;
}

//
// PROXY_RESPONSE
//
HttpProxyInstance::proxy_state HttpProxyInstance::process_response()
{
    new_packet();

    packet().source(socket_get_peer_address(_server->get_fd()));
    packet().destination(socket_get_peer_address(_client->get_fd()));

    if (!packet().headline().read_in(_server) || 
	!packet().headers().read_in(_server))
    {
	return PROXY_ERROR;
    }

    if (packet().headline().full().find("HTTP")!=0)
    {	
	logger.message(logger.LOG_ERROR,"Malformed response header: %s",packet().headline().full().c_str());
	return PROXY_ERROR;
    }

    logger.message(logger.LOG_DEBUG,"Processing response %s",packet().headline().full().c_str());
    packet().headers().set("Connection","close");  
    packet().log();

    if (!_proxy->filter_packet(packet())) return PROXY_FILTER_BLOCK; 

    // FIXME: process HEAD requests responses differently.

    if (!packet().headline().write_out(_client) ||
    	!packet().headers().write_out(_client) ||
   	!packet().content_read_in_write_out(_server,_client) ||
	!_client->write_flush())
    {
	logger.message(logger.LOG_ERROR,"Failed to send response packet to client");
	return PROXY_ERROR;
    }
    logger.message(logger.LOG_DEBUG,"Response sent back to client");
    
    return PROXY_CONNECTION_CLOSE;
}

//
// PROXY_RESPONSE_NOT_FOUND
//
HttpProxyInstance::proxy_state HttpProxyInstance::process_response_not_found()
{
    new_packet();

    packet().destination(socket_get_local_address(_client->get_fd()));
    packet().source(socket_get_peer_address(_client->get_fd()));

    packet().headline().set("HTTP/1.1 404 Not Found\r\n");
    packet().headers().set("Proxy-agent", "Proxy Certs");
    packet().headers().set("Content-type", "text/html");
    packet().headers().set("Connection", "close");
    packet().content_add("<html><body><h1>Not found by proxy</h1></body></html>");
    packet().headers().set("Content-length", packet().content_length()); 
    
    logger.message(logger.LOG_DEBUG,"Sending response %s",packet().headline().full().c_str());
    
    packet().log();	

    if (!packet().headline().write_out(_client) ||
	!packet().headers().write_out(_client) || 
	!packet().content_write_out(_client) ||
	!_client->write_flush())
    {
	logger.message(logger.LOG_ERROR,"Failed to send response packet to client");
	return PROXY_ERROR;
    }
    logger.message(logger.LOG_DEBUG,"Response sent back to client");

    return PROXY_CONNECTION_CLOSE;
}

//
// PROXY_RESPONSE_GATEWAY_TIMEOUT
//
HttpProxyInstance::proxy_state HttpProxyInstance::process_response_gateway_timeout()
{
    new_packet();
	
    packet().destination(socket_get_local_address(_client->get_fd()));
    packet().source(socket_get_peer_address(_client->get_fd()));

    packet().headline().set("HTTP/1.1 504 Gateway timeout\r\n");
    packet().headers().set("Proxy-agent","Proxy Certs");
    packet().headers().set("Content-type","text/html");
    packet().headers().set("Connection", "close");
    packet().content_add("<html><body><h1>Gateway timeout from proxy</h1></body></html>");
    packet().headers().set("Content-length",packet().content_length()); 

    logger.message(logger.LOG_DEBUG,"Sending response %s",packet().headline().full().c_str());
    
    packet().log();	

    if (!packet().headline().write_out(_client) ||
	!packet().headers().write_out(_client) || 
	!packet().content_write_out(_client) ||
	!_client->write_flush())
    {
	logger.message(logger.LOG_ERROR,"Failed to send response packet to client");
	return PROXY_ERROR;
    } 
    logger.message(logger.LOG_DEBUG,"Response sent back to client");
  
    return PROXY_CONNECTION_CLOSE;
}

//
// PROXY_RESPONSE_INTERNAL
//
HttpProxyInstance::proxy_state HttpProxyInstance::process_response_internal()
{
    new_packet();

    packet().destination(socket_get_local_address(_client->get_fd()));
    packet().source(socket_get_peer_address(_client->get_fd()));

    packet().headline().set("HTTP/1.1 200 OK\r\n");
    packet().headers().set("Proxy-agent","Proxy Certs");
    packet().headers().set("Content-type","text/html");
    packet().headers().set("Connection","close");
    packet().content_add("<html><body><h1>Welcome to the proxy</h1></body></html>");
    packet().headers().set("Content-length",packet().content_length()); 

    logger.message(logger.LOG_DEBUG,"Sending response %s",packet().headline().full().c_str());
    
    packet().log();	

    if (!packet().headline().write_out(_client) ||
	!packet().headers().write_out(_client) || 
	!packet().content_write_out(_client) || 
	!_client->write_flush())
    {
	logger.message(logger.LOG_ERROR,"Failed to send response packet to client");
	return PROXY_ERROR;
    } 
    logger.message(logger.LOG_DEBUG,"Response sent back to client");

    return PROXY_CONNECTION_CLOSE;  
}


bool HttpProxyInstance::run()
{
    HttpPacketDB db;
    bool running = true;
    proxy_state state = PROXY_INIT;

    if (_client==NULL)
	return false;

    logger.message(logger.LOG_DEBUG,"Running proxy with direct connection");

    while (running) 
    {
	switch (state) {
	    case PROXY_INIT:
	    case PROXY_REQUEST:
	        state = process_request();
		break;
	    case PROXY_RESPONSE:
		state = process_response();
		break;
	    case PROXY_RESPONSE_GATEWAY_TIMEOUT:
		state = process_response_gateway_timeout();
		break;
	    case PROXY_RESPONSE_NOT_FOUND:
		state = process_response_not_found();
		break;
	    case PROXY_TUNNEL:
		state = process_tunnel_setup();
		break;
	    case PROXY_SSL_INIT:
	    case PROXY_SSL_REQUEST:
		state = process_ssl_request();
		break;
	    case PROXY_SSL_RESPONSE:
		state = process_ssl_response();
		break;
	    case PROXY_RESPONSE_INTERNAL:
		state = process_response_internal();
		break;
	    case PROXY_FILTER_BLOCK:
	    case PROXY_CONNECTION_CLOSE:
	    case PROXY_ERROR:
		running = false;
		break;
	}

	if (has_packet())
	{ 
		if (running && _save_packet)
		{
			db.open(false);
			db.store(_current_packet);
			db.close();
		}
	}

    }

    if (state==PROXY_CONNECTION_CLOSE)
	logger.message(logger.LOG_DEBUG,"Connection ended through 'Connection: close' header");
    else if (_client && _client->read_end())
	logger.message(logger.LOG_DEBUG,"Connection ended by client");
    else if (_server && _server->read_end())
	logger.message(logger.LOG_DEBUG,"Connection ended by server");
    else
	logger.message(logger.LOG_DEBUG,"Connection ended because of error");
//    return _client->eof() || _client->eof() || _connection_close;
    return true;
}

/************************/

